from data_load import dataManager as dm
from FLServer import FLServer

# 为了保证数据规模与集中训练一直，保证client_number*data_counts=60000 

dataset_params={
    # 客户端数量
    "client_number":20,
    # 客户端数据规模
    "data_counts":30,
    # 单个客户端类别舒朗
    "class_counts":2,
}

server_params = {
    # 迭代的轮次
    "rounds":100,
    # 客户端参与度。表示每轮有多少个客户端参与到训练当中
    "ratio":1.0,
}

client_params = {
    # prox本地更新的权重
    "prox_mu": 0.1,
    # 迭代轮次
    "epochs" :1,
    # batchsize
    "batch_size":128
}

# 初始化各种类型的数据集。这里返回的是不同数据的索引，不能修改，否则会影响原来的数据。
train_dataset_list_avg = dm.allocate_data_avg(client_number=dataset_params["client_number"])
train_dataset_list_iid = dm.allocate_data_iid(client_number=dataset_params["client_number"],data_counts=dataset_params["data_counts"])
train_dataset_list_noniid = dm.allocate_data_noniid(client_number=dataset_params["client_number"],class_counts=dataset_params["class_counts"],data_counts=dataset_params["data_counts"])
test_dataset_list = dm.allocate_data_test()

# ------------------------------------------------fedavg----------------------------------------
# # 初始化，并开始训练。就这样了。init作为初始话的内容。传递进入dataset，其他的参数都是可选的的
# server = FLServer(train_dataset_list_avg,test_dataset_list,server_params,client_params)
# server.start_federated()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_iid,test_dataset_list,server_params,client_params)
# server.start_federated()

# 初始化并开始训练
# server = FLServer(train_dataset_list_noniid,test_dataset_list,server_params,client_params)
# server.start_federated() 


# ------------------------------------------------fedida----------------------------------------
# 初始化，并开始训练
# server = FLServer(train_dataset_list_avg,test_dataset_list,server_params,client_params)
# server.start_federated_ida()

# 初始化并开始训练
# server = FLServer(train_dataset_list_iid,test_dataset_list,server_params,client_params)
# server.start_federated_ida()

# 初始化并开始训练
# server = FLServer(train_dataset_list_noniid,test_dataset_list,server_params,client_params)
# server.start_federated_ida() 


# --------------------------------------------------fedprox----------------------------------------
# 初始化，并开始训练
# server = FLServer(train_dataset_list_avg,test_dataset_list,server_params,client_params)
# server.start_federated_prox()

# 初始化并开始训练
# server = FLServer(train_dataset_list_iid,test_dataset_list,server_params,client_params)
# server.start_federated_prox()

# 初始化并开始训练
# server = FLServer(train_dataset_list_noniid,test_dataset_list,server_params,client_params)
# server.start_federated_prox() 


# -------------------------------------------------fedper---------------------------------------------

# # 初始化，并开始训练
# server = FLServer(train_dataset_list_avg,test_dataset_list,server_params,client_params)
# server.start_federated_meta()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_iid,test_dataset_list,server_params,client_params)
# server.start_federated_meta()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_noniid,test_dataset_list,server_params,client_params)
# server.start_federated_meta() 



# -------------------------------------------------fedmaml---------------------------------------------

# # 初始化，并开始训练
# server = FLServer(train_dataset_list_avg,test_dataset_list,server_params,client_params)
# server.start_federated_maml()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_iid,test_dataset_list,server_params,client_params)
# server.start_federated_maml()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_noniid,test_dataset_list,server_params,client_params)
# server.start_federated_maml() 


# -------------------------------------------------fedmaml_by_grad---------------------------------------------

# 初始化，并开始训练
# server = FLServer(train_dataset_list_avg,test_dataset_list,server_params,client_params)
# server.start_federated_maml_by_grad()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_iid,test_dataset_list,server_params,client_params)
# server.start_federated_maml_by_grad()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_noniid,test_dataset_list,server_params,client_params)
# server.start_federated_maml_by_grad() 




# -------------------------------------------------fedreptile---------------------------------------------

# # 初始化，并开始训练
# server = FLServer(train_dataset_list_avg,test_dataset_list,server_params,client_params)
# server.start_federated_reptile()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_iid,test_dataset_list,server_params,client_params)
# server.start_federated_reptile()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_noniid,test_dataset_list,server_params,client_params)
# server.start_federated_reptile() 


# -------------------------------------------------fedreptile_by_grad---------------------------------------------

# # 初始化，并开始训练
# server = FLServer(train_dataset_list_avg,test_dataset_list,server_params,client_params)
# server.start_federated_reptile_by_grad()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_iid,test_dataset_list,server_params,client_params)
# server.start_federated_reptile_by_grad()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_noniid,test_dataset_list,server_params,client_params)
# server.start_federated_reptile_by_grad() 


# -------------------------------------------------fedper---------------------------------------------

# # 初始化，并开始训练
# server = FLServer(train_dataset_list_avg,test_dataset_list,server_params,client_params)
# server.start_federated_per()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_iid,test_dataset_list,server_params,client_params)
# server.start_federated_per()

# # 初始化并开始训练
# server = FLServer(train_dataset_list_noniid,test_dataset_list,server_params,client_params)
# server.start_federated_per() 
